from tree_sitter import Language, Parser, Node
from init_parser import *
from entity import If_node, TRUE_TO_FALSE
from util import merge_var_change, format_compare
from z3 import Real, Solver, sat, And, Or, Not, simplify, Int
import re
from copy import deepcopy
from typing import Dict, List  # 导入 Dict

from graphviz import Digraph
import os


# 2、ast中提取 if 语句结构，并处理变量的赋值，将结果储存在if_node
def format_tree_cpp(root: Node, if_tree: If_node, vars: Dict[str, str]):
    # 这个方法将一个语法树中的if语句抽象出来构造if树,只关注if语句,先忽略if语句块中的其他语句比如赋值和运算
    for child in root.children:
        # 这一步是查询出所有在if之前的对变量的修改
        # 判断当前节点是否只有一个子节点，如果是则可能是赋值语句
        # print('')
        # print('语句子节点len:', child.children.__len__())
        # print('语句类型:', child.type)

        # 处理赋值语句
        if child.children.__len__() == 2:
            # 将当前节点的唯一子节点赋值给child1（赋值语句）
            child1 = child.children[0]
            # 检查当前节点是否为赋值语句
            if child1.type == 'assignment_expression':
                # 赋值语句的左值（变量名）和右值（值/表达式）分别赋值
                l_value = str(child1.children[0].text, encoding='utf-8')
                r_value = str(child1.children[2].text, encoding='utf-8')
                # print('赋值语句左：', l_value)
                # print('赋值语句左：', r_value)
                # 如果赋值的右值是调用函数,说明赋值是通过函数调用完成的-不处理
                if child1.children[2].type == 'call':
                    # l_value暂时是它自身的符号
                    vars[l_value] = l_value
                    # print(vars[l_value])
                else:
                    # 右值不是函数
                    # 如果l_value存在于词典，则之前该变量被赋值过，我们需要更新对照表
                    if l_value in vars:
                        vars[l_value] = vars[l_value].replace(l_value, f'({r_value})')
                        # print('l_value已存在于词典,reset_var:', vars[l_value], vars)
                    else:
                        # 如果l_value不在词典，则是第一次被赋值，变量名就是符号对照表对应的符号
                        vars[l_value] = r_value
                        # print('new_var:', vars, ' ', vars[l_value])

        # 处理if语句
        if child.type == 'if_statement':
            # 获取if语句节点的所有子节点
            cc = child.children

            i = 0
            for c in cc:
                # print('c.type', i, ':', c.type, 'c.text: ', str(c.text, encoding='utf-8'))
                i += 1

            # 将一个新的if_node节点加入到if语法树
            if_tree.children.append(If_node('if', merge_var_change(str(cc[1].text, encoding='utf-8'), vars)))
            print('if_tree')
            print_tree_cpp(if_tree)
            # 递归调用，下一步为if的block块，if_tree当前节点为刚刚加入的子节点
            # deepcopy(vars) 用于传递 vars 字典的副本，以避免在递归过程中修改原始字典
            # 递归if中的语句？
            format_tree_cpp(cc[2], if_tree.children[-1], deepcopy(vars))

            # 如果有else和else if，那么必定在3项以后
            i = 0
            for c in cc[3:]:
                # print('else和else if', i, ': ', c.type)
                j = 0
                for ccc in c.children:
                    # print('c.child', j, ':', ccc.type, str(ccc.text, encoding='utf-8'))
                    j += 1

                # 处理 else if 子句
                if c.children[0].type == 'else' and c.children[1].type == 'if_statement':
                    # print('else if')
                    # 获取 else if 子句的 if_statement 节点
                    else_if_statement = c.children[1]
                    # 解析 else if 子句中的条件表达式，并将其添加到 if_tree 中
                    condition_expression = str(else_if_statement.children[1].text, encoding='utf-8')
                    if_tree.children.append(If_node('else if', merge_var_change(condition_expression, vars)))
                    print_tree_cpp(if_tree)
                    # 递归处理 else if 子句的语句块
                    format_tree_cpp(else_if_statement.children[2], if_tree.children[-1], deepcopy(vars))

                # 处理 else 子句
                if c.children[0].type == 'else':
                    # print('else')
                    if_tree.children.append(If_node('else', None))
                    print_tree_cpp(if_tree)
                    format_tree_cpp(c.children[0], if_tree.children[-1], deepcopy(vars))


# 3、 生成路径覆盖测试用例的核心函数，为每个 if 节点生成所有可能的路径约束条件。它递归地处理每条路径，并生成相应的约束条件。
def get_truth_trees_cpp(root: If_node):
    '''
    使用路径覆盖
    '''

    # 考虑嵌套
    # 路径覆盖可以直接获得约束条件
    # 递归地生成从当前节点开始的路径上的约束条件
    def get_re(root: If_node, prefix: list, res: List[List[str]]):
        # 如果当前比较条件不是none
        if root.compare_statement != None:
            # 格式化条件并加入到prefix
            prefix.append(format_compare(root.compare_statement))
        i = root.children.__len__()
        # 当前节点没有子节点
        if i == 0:
            if prefix != []:
                # 将prefix加入到res
                res.append(prefix)
            return
        # 遍历当前节点的所有子节点
        for ii in range(i):
            # 从当前节点的第ii个子节点递归调用函数，生成从该节点开始的路径约束
            # deepcopy(prefix)是对prefix深拷贝
            get_re(root.children[ii], deepcopy(prefix), res)
            if root.children[ii].compare_statement != None:
                # 将当前子节点的条件的否定添加到prefix
                prefix.append(f'Not({format_compare(root.children[ii].compare_statement)})')

    l = len(root.children)
    res: List[List[str]] = []
    prefix = []  # 到达这个分支需要的前置条件
    for i in range(l):
        get_re(root.children[i], deepcopy(prefix), res)
        if root.children[i].compare_statement != None:
            prefix.append(f'Not({format_compare(root.children[i].compare_statement)})')
    return res


# 4、打印 if 树结构，用于调试或可视化。它递归地遍历 If_node 树，并打印每个节点的详细信息
def print_tree_cpp(root: If_node, level=0):
    '''
    一个打印if树的方法
    '''
    # level表示当前节点在树中的层级  repr（root）返回root的字符串表示
    print('  ' * level + repr(root))
    for child in root.children:
        print_tree_cpp(child, level + 1)
        # module(None)
        # if ((a + 1) > 9)
        #     if (((a + 1) + 1) > 10)
        # elif ((b + 2) >= (a + 1) + 5)
        # else (None)


# 5、在字符串中替换变量名，根据给定的字典，将变量名替换为实际的值。
def replace_keys_with_values_cpp(string: str, rep_dict: dict):
    if string is None:
        return None
    for key, value in rep_dict.items():
        if key in string:
            string = string.replace(key, value)
            break
    return string


# 6、从字符串中提取变量名，用于提取约束条件中的变量，排除掉某些关键字。
def get_variables_cpp(s: str):
    variables = re.findall(r"\b[a-zA-Z]+\b", s)
    variables = list(filter(lambda x: x not in ['And', 'Not', 'Or'], variables))
    return variables


# 选择 C++ 解析语言
CPP_LANGUAGE = Language('build/my-languages.so', 'cpp')
cpp_parser = Parser()
# parser的一个实例，设置解析语言为cpp
cpp_parser.set_language(CPP_LANGUAGE)


# 按钮-生成测试用例
# 1、入口，调用其他函数来解析代码、构建 if 树、生成约束条件，并最终输出测试用例。
def analyze_cpp(code: str):
    # ----------------------生成中间表示（抽象语法树）------------------------------
    result = ''
    # 解析python代码为抽象语法树，将字符串转换为byte，并将byte流输入到解析器，生成抽象语法树
    tree = cpp_parser.parse(bytes(code, "utf8"))
    # ----------------------进行符号执行-----------------------------------------
    # tree.root_node抽象语法树的根节点
    root_node = tree.root_node
    # If_node是自定义的类（在entity.py）中，表示if语句节点。下面是在构建if语法树的根节点
    if_tree = If_node('module', None)
    # 下面方法从抽象语法树根节点提取if语句相关结构，并将信息存到if语法树。
    # {}是空字典 作用是在符号执行遍历if语法树时更新符号变量对照表
    format_tree_cpp(root_node, if_tree, {})
    # ------------------------生成约束条件----------------------------------------
    # 打印if树
    # print_tree(if_tree)
    # 生成路径约束，一条路径的约束条件是一个数组
    restrictions = get_truth_trees_cpp(if_tree)
    # # 保存 if 树的图像
    # images_dir = 'static/images'
    # save_ast_image(if_tree, images_dir)
    # print(restrictions)
    # --------------------------构造Z3求解器求解约束--------------------------------
    letters = set()
    # 遍历约束条件列表
    for res in restrictions:
        # 遍历约束条件
        for r in res:
            # 从约束条件中提取变量名
            v = get_variables_cpp(r)
            # 变量名加到letters中
            letters.update(v)
    # print(letters)
    # 创建字典，包含Z3的逻辑操作符
    variables = {'And': And, 'Or': Or, 'Not': Not}
    # 遍历变量名
    for x in letters:
        # 为每一个变量x创建一个整型变量添加到字典
        variables[x] = Int(x)
    result = result + '一组路径覆盖测试用例如下\n'
    # 遍历约束列表
    for res in restrictions:
        # 针对一条路径的约束
        # 创建Z3求解器的实例
        s = Solver()
        # 针对当前约束列表中的一个约束
        for r in res:
            # eval（）将字符串形式的r转换为Z3表达式
            # simplify简化Z3表达式
            rr = simplify(eval(r, {}, variables))
            # print(rr)
            # 将简化后的约束条件添加到Z3求解器
            s.add(rr)
        result = result + f'该测试用例约束条件为:{res}\n'
        # 检查当前Z3求解器中的约束是否可满足，可满足返回sat
        if s.check() == sat:
            # 获取一个解
            m = s.model()
            result = result + str(m) + '\n'
        else:
            result = result + "无解\n"
    return result


if __name__ == "__main__":
    code = """
        int a, b;
        cin >> a;
        cin >> b;
        a = a + 1;
        b = b + 2;
        if (a > 5) {
            a=a+1;
            if (a > 10){
                cout << "11";
            }
        }
        else if (b >= a + 3) {
            cout << "22";
        }
        else {
            cout << "33";
        }
        return 0;
    """
    print(analyze_cpp(code))

    # code = """
    #     int a, b;
    #     cin >> a;
    #     cin >> b;
    #     a = a + 1;
    #     b = b + 2;
    #     if (a > 9) {
    #         a=a+1;
    #         if (a > 10){
    #             cout << "11";
    #         }
    #     }else{
    #         cout << "33";
    #     }
    #     return 0;
    # """
